import json

import pandas as pd
from langchain_community.chat_models import ChatZhipuAI


def get_dict():
    with open("dict.json", "r", encoding="utf-8") as file:
        return json.load(file)


dict_data = get_dict()


def extract_table_content(response):
    tables = dict_data
    for table in tables:
        if response == table["数据表名"]:
            return table
    return None


def choose_table(question):
    llm = ChatZhipuAI(
        model="glm-4-plus",
        zhipuai_api_key="dae00bbfa5604ca49303bdc0ed54b95f.KCE7ukYd34wWi0AI"
    )
    context_text = str(dict_data)
    prompt = f"""我有如下数据表：<{context_text}>
    现在基于数据表回答问题：{question}。
    分析需要哪些数据表。
    仅返回需要的数据表名，无需展示分析过程，也不要添加额外的标点符号。
    """
    res = llm.invoke(prompt)

    return res


def load_and_filter_data(file_path, start_time, end_time, power_column):
    """
    加载 CSV 文件并筛选指定时间范围内的数据
    :param file_path: CSV 文件路径
    :param start_time: 开始时间
    :param end_time: 结束时间
    :param power_column: 功率列名
    :return: 筛选后的 DataFrame
    """
    try:
        df = pd.read_csv(file_path)
    except FileNotFoundError:
        raise FileNotFoundError(f"文件 {file_path} 未找到")

    # 确保时间列是 datetime 类型
    try:
        df["csvTime"] = pd.to_datetime(df["csvTime"])
    except Exception as e:
        raise ValueError(f"时间列转换失败: {e}")

    # 筛选特定时间范围内的数据
    filtered_data = df[
        (df["csvTime"] >= start_time) & (df["csvTime"] < end_time)
        ].copy()

    if filtered_data.empty:
        return None

    # 计算时间差（秒）
    filtered_data.loc[:, "diff_seconds"] = (
        filtered_data["csvTime"].diff().dt.total_seconds().shift(-1)
    )

    # 计算每个时间间隔的能耗（kWh）
    filtered_data.loc[:, "energy_kWh"] = (
            filtered_data["diff_seconds"] * filtered_data[power_column] / 3600
    )

    return filtered_data

def calculate_total_energy(start_time, end_time, device_name="折臂吊车"):
    """
    计算指定时间段内的总能耗
    :param start_time: 查询的开始时间（字符串或 datetime 类型）
    :param end_time: 查询的结束时间（字符串或 datetime 类型）
    :param device_name: 设备名称，默认为 '折臂吊车'
    :return: 总能耗（kWh，float 类型）
    """
    # 设备配置映射：设备名称 -> (表名, 功率列名)
    device_config = {
        "折臂吊车": ("device_13_11_meter_1311", "13-11-6_v"),
        "一号门架": ("device_1_5_meter_105", "1-5-6_v"),  # 一号门架的配置
        "二号门架": ("device_13_14_meter_1314", "13-14-6_v"),  # 二号门架的配置
        "绞车": ("device_1_15_meter_115", "1-15-6_v"),  # 添加绞车的配置
    }

    # 检查设备名称是否有效
    if device_name not in device_config:
        raise ValueError(f"未知的设备名称: {device_name}")

    # 获取设备配置
    table_name, power_column = device_config[device_name]

    # 读取 CSV 文件并计算能耗
    file_path = f"database_in_use/{table_name}.csv"
    try:
        filtered_data = load_and_filter_data(
            file_path, start_time, end_time, power_column
        )
        if filtered_data is None:
            return None
        total_energy_kWh = filtered_data["energy_kWh"].sum()
        return round(abs(total_energy_kWh), 2)
    except Exception as e:
        raise ValueError(f"计算能耗时出错: {e}")
